"""
文档处理和向量化服务（含OCR支持）
"""
from typing import List, Dict, Optional
import logging
import os
import re
from pathlib import Path
from app.services.ocr_service import ocr_service

logger = logging.getLogger(__name__)


class DocumentProcessor:
    """文档处理器 - 负责文本提取和分块"""
    
    def __init__(self, chunk_size: int = 500, chunk_overlap: int = 50):
        """
        初始化文档处理器
        
        Args:
            chunk_size: 文本块大小（字符数）
            chunk_overlap: 文本块重叠大小
        """
        self.chunk_size = chunk_size
        self.chunk_overlap = chunk_overlap
    
    def extract_text(self, file_path: str, file_type: str) -> str:
        """
        从文件中提取文本（支持OCR）
        
        Args:
            file_path: 文件路径
            file_type: 文件MIME类型
            
        Returns:
            提取的文本内容
        """
        try:
            # 检查是否为图片文件
            if 'image' in file_type or ocr_service.is_image_file(file_path):
                logger.info(f"检测到图片文件，使用OCR提取文本: {file_path}")
                return self._extract_image_ocr(file_path)
            
            if file_type.startswith('text/'):
                # 纯文本文件
                return self._extract_text_file(file_path)
            elif 'pdf' in file_type:
                # PDF文件（先尝试文本提取，失败则OCR）
                return self._extract_pdf(file_path)
            elif 'word' in file_type or file_path.endswith(('.doc', '.docx')):
                # Word文档
                return self._extract_word(file_path)
            elif 'excel' in file_type or 'spreadsheet' in file_type or file_path.endswith(('.xls', '.xlsx')):
                # Excel文件
                return self._extract_excel(file_path)
            elif 'presentation' in file_type or file_path.endswith(('.ppt', '.pptx')):
                # PowerPoint文件
                return self._extract_ppt(file_path)
            else:
                logger.warning(f"不支持的文件类型: {file_type}")
                return ""
        except Exception as e:
            logger.error(f"提取文本失败: {str(e)}")
            raise
    
    def _extract_text_file(self, file_path: str) -> str:
        """提取纯文本文件"""
        try:
            with open(file_path, 'r', encoding='utf-8') as f:
                return f.read()
        except UnicodeDecodeError:
            # 尝试其他编码
            with open(file_path, 'r', encoding='gbk') as f:
                return f.read()
    
    def _extract_pdf(self, file_path: str) -> str:
        """提取PDF文本（含OCR备用）"""
        try:
            import pypdf
            text = []
            with open(file_path, 'rb') as f:
                pdf_reader = pypdf.PdfReader(f)
                for page in pdf_reader.pages:
                    page_text = page.extract_text()
                    text.append(page_text)
            
            full_text = '\n'.join(text)
            
            # 如果提取的文本太少（可能是扫描件），使用OCR
            if len(full_text.strip()) < 100:
                logger.info(f"PDF文本提取结果太少（{len(full_text)}字符），尝试OCR识别...")
                try:
                    ocr_text = ocr_service.extract_text_from_pdf_images(file_path)
                    if len(ocr_text) > len(full_text):
                        logger.info("✅ OCR识别效果更好，使用OCR结果")
                        return ocr_text
                except Exception as ocr_error:
                    logger.warning(f"⚠️ OCR识别失败，使用原始文本: {str(ocr_error)}")
            
            return full_text
            
        except ImportError:
            logger.warning("pypdf未安装，无法处理PDF文件")
            return ""
        except Exception as e:
            logger.error(f"提取PDF失败: {str(e)}")
            return ""
    
    def _extract_word(self, file_path: str) -> str:
        """提取Word文档文本"""
        try:
            import docx
            doc = docx.Document(file_path)
            text = []
            for paragraph in doc.paragraphs:
                text.append(paragraph.text)
            return '\n'.join(text)
        except ImportError:
            logger.warning("python-docx未安装，无法处理Word文件")
            return ""
        except Exception as e:
            logger.error(f"提取Word失败: {str(e)}")
            return ""
    
    def _extract_image_ocr(self, file_path: str) -> str:
        """使用OCR提取图片文本"""
        try:
            text, details = ocr_service.extract_text_from_image(file_path)
            
            # 记录识别信息
            logger.info(f"✅ OCR识别完成: {len(details)} 行文字, 共 {len(text)} 字符")
            
            return text
        except Exception as e:
            logger.error(f"❌ OCR识别失败: {str(e)}")
            return ""
    
    def _extract_excel(self, file_path: str) -> str:
        """提取Excel文本"""
        try:
            import openpyxl
            wb = openpyxl.load_workbook(file_path, read_only=True)
            text = []
            for sheet in wb.worksheets:
                text.append(f"工作表: {sheet.title}")
                for row in sheet.iter_rows(values_only=True):
                    row_text = '\t'.join([str(cell) if cell is not None else '' for cell in row])
                    if row_text.strip():
                        text.append(row_text)
            return '\n'.join(text)
        except ImportError:
            logger.warning("openpyxl未安装，无法处理Excel文件")
            return ""
        except Exception as e:
            logger.error(f"提取Excel失败: {str(e)}")
            return ""
    
    def _extract_ppt(self, file_path: str) -> str:
        """提取PowerPoint文本"""
        try:
            from pptx import Presentation
            prs = Presentation(file_path)
            text = []
            
            for i, slide in enumerate(prs.slides, 1):
                slide_text = [f"--- 幻灯片 {i} ---"]
                
                # 提取标题和内容
                for shape in slide.shapes:
                    if hasattr(shape, "text") and shape.text.strip():
                        slide_text.append(shape.text.strip())
                    
                    # 提取表格内容
                    if shape.has_table:
                        table = shape.table
                        for row in table.rows:
                            row_text = '\t'.join([cell.text.strip() for cell in row.cells if cell.text.strip()])
                            if row_text:
                                slide_text.append(row_text)
                
                if len(slide_text) > 1:  # 如果有内容（除了标题）
                    text.append('\n'.join(slide_text))
            
            return '\n\n'.join(text)
        except ImportError:
            logger.warning("python-pptx未安装，无法处理PowerPoint文件")
            return ""
        except Exception as e:
            logger.error(f"提取PowerPoint失败: {str(e)}")
            return ""
    
    def split_text(self, text: str) -> List[str]:
        """
        将文本分割成块
        
        Args:
            text: 原始文本
            
        Returns:
            文本块列表
        """
        if not text or len(text) == 0:
            return []
        
        # 清理文本
        text = self._clean_text(text)
        
        # 按段落分割
        paragraphs = re.split(r'\n\s*\n', text)
        
        chunks = []
        current_chunk = ""
        
        for paragraph in paragraphs:
            paragraph = paragraph.strip()
            if not paragraph:
                continue
            
            # 如果当前块加上新段落不超过块大小，则添加
            if len(current_chunk) + len(paragraph) + 2 <= self.chunk_size:
                if current_chunk:
                    current_chunk += "\n\n" + paragraph
                else:
                    current_chunk = paragraph
            else:
                # 否则保存当前块，开始新块
                if current_chunk:
                    chunks.append(current_chunk)
                
                # 如果段落本身超过块大小，需要进一步分割
                if len(paragraph) > self.chunk_size:
                    # 按句子分割
                    sentences = re.split(r'[。！？\n]', paragraph)
                    temp_chunk = ""
                    for sentence in sentences:
                        if not sentence.strip():
                            continue
                        if len(temp_chunk) + len(sentence) + 1 <= self.chunk_size:
                            temp_chunk += sentence + "。"
                        else:
                            if temp_chunk:
                                chunks.append(temp_chunk)
                            temp_chunk = sentence + "。"
                    if temp_chunk:
                        current_chunk = temp_chunk
                    else:
                        current_chunk = ""
                else:
                    current_chunk = paragraph
        
        # 添加最后一块
        if current_chunk:
            chunks.append(current_chunk)
        
        return chunks
    
    def _clean_text(self, text: str) -> str:
        """清理文本"""
        # 删除多余空白
        text = re.sub(r'\s+', ' ', text)
        # 恢复段落分隔
        text = re.sub(r'\s*\n\s*', '\n', text)
        return text.strip()
    
    def process_document(self, file_path: str, file_type: str) -> Dict:
        """
        处理文档：提取文本并分块
        
        Args:
            file_path: 文件路径
            file_type: 文件类型
            
        Returns:
            包含文本块和统计信息的字典
        """
        try:
            # 提取文本
            text = self.extract_text(file_path, file_type)
            
            if not text:
                return {
                    "chunks": [],
                    "word_count": 0,
                    "chunk_count": 0,
                    "summary": "无法提取文本内容"
                }
            
            # 分块
            chunks = self.split_text(text)
            
            # 统计
            word_count = len(text)
            chunk_count = len(chunks)
            
            # 生成摘要（取前500字符，更完整）
            summary = text[:500] + "..." if len(text) > 500 else text
            
            return {
                "chunks": chunks,
                "word_count": word_count,
                "chunk_count": chunk_count,
                "summary": summary
            }
        except Exception as e:
            logger.error(f"处理文档失败: {str(e)}")
            raise


# 全局文档处理器实例
document_processor = DocumentProcessor()

